package com.linkhub.oms.component.trace.core;

import com.google.common.collect.Maps;
import com.linkhub.oms.common.util.json.QtFastjsonUtil;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Description: SpanId
 * @Author: echo
 * @Date: 2022/8/31 17:27
 * @Version: 1.0
 */
@Slf4j
public class LeafSpanIdInfo {
    /**
     * 记录基于 {@code traceId_parentSpanId} 为key, 的 {@code spanId} 数量
     *
     */
    private final Map<String, AtomicInteger> traceIdCounter = Maps.newConcurrentMap();

    /**
     * 获取并递增 {@code spanId}, 该方法必须保证线程安全
     *
     * @param traceIdAndParentSpanId traceIdAndParentSpanId
     * @return spanId
     */
    public final int nextSpanId(String traceIdAndParentSpanId) {
        // 相同key只创建一个AtomicInteger
        log.debug("traceIdAndParentSpanId = [{}], traceIdCounter = {}", traceIdAndParentSpanId, QtFastjsonUtil.toJSONString(traceIdCounter));
        AtomicInteger count = create(traceIdAndParentSpanId);
        int result = count.incrementAndGet();
        traceIdCounter.putIfAbsent(traceIdAndParentSpanId, count);
        return result;

    }

    public int get(String traceIdAndParentSpanId) {
        AtomicInteger count = create(traceIdAndParentSpanId);
        int result = count.get();
        return result == 0 ? 1 : result;
    }

    private AtomicInteger create(String traceIdAndParentSpanId) {
        return traceIdCounter.computeIfAbsent(traceIdAndParentSpanId, key -> new AtomicInteger(0));
    }

}
